#%%
import numpy as np, matplotlib.pyplot as plt, seaborn as sns, pandas as pd
%matplotlib inline

grains = ['dn', 'jo', 'mb', 'mk', 'pg', 'sf']
labels_old = {'dn': 'dragon glass', 'jo': 'joels gravel', 'mb': 'marble chips', 'mk': 'michaels gravel', 'pg': 'pea gravel', 'sf': 'shell fragments'}
labels_new = {'dn': 'Tempered glass chips', 'jo': 'Natural gravel 2', 'mb': 'Marble chips', 'mk': 'Natural gravel 3', 'pg': 'Pea gravel', 'sf': 'Shell fragments'}


# %%
# load up grain density data
df_density = pd.read_csv('density_and_volume/2022_grain_density_and_volume_measurements_simple.csv').set_index('grain')

rho = {grain: df_density.grains_rho[grain] for grain in grains}
# drho = 
d_VES = {grain: df_density.D_VES[grain]/1000 for grain in grains}
# dd_VES = 

# %%
# load up and process angle of repose data
df_aor = {}
for grain in grains:
    aor = pd.read_csv('AOR/%s/%s_AoR.csv'%(grain, grain)).values[0]
    aor_arr = np.asarray(aor[1:]).astype(float)
    n = np.isfinite(aor_arr).sum()
    df_aor[grain] = {'mean': np.nanmean(aor_arr), 'se': np.nanstd(aor_arr)/np.sqrt(n)}
df_aor = pd.DataFrame(df_aor).T
aor = {grain: df_aor['mean'][grain] for grain in grains}
daor = {grain: df_aor.se[grain] for grain in grains}

# %%
# load up settling velocity data and process it
df_vel = {}
for grain in grains:
    vels = {}
    if (grain is 'dn') or (grain is 'mb') or (grain is 'pg'): date = '20220208'
    else: date = '20220210'

    vels = pd.read_csv('settling_velocity/%s/%s/%s_%i_settling_velocities.csv'%(date, grain, grain, 1))
    for i in range(2,5):
        try:
            temp = pd.read_csv('settling_velocity/%s/%s/%s_%i_settling_velocities.csv'%(date, grain, grain, i))
        except FileNotFoundError: 
            continue
        vels = pd.concat([vels, temp])

    vels = vels.reset_index().drop('index', axis=1).drop('particleId', axis=1)
    df_vel[grain] = {'mean': vels.vel.mean()/1000, 'se': vels.vel.std()/1000/np.sqrt(vels.vel.values.size)}
df_vel = pd.DataFrame(df_vel).T
vel = {grain: df_vel['mean'][grain] for grain in grains}
dvel = {grain: df_vel.se[grain] for grain in grains}

# %%
# load up VES settling velocity predictions
df_ves_taylor = pd.read_csv('5_var_VES_settling_velocity.csv').set_index('Unnamed: 0').rename_axis(None)
vel_VES = {grain: df_ves_taylor.ws_sphere[grain] for grain in grains}
dvel_VES = {grain: df_ves_taylor.dws_sphere[grain] for grain in grains}

# %%
# process CSF and shape data
# CSF_taylor = {'dn': 0.721, 'jo': 0.674, 'mb': 0.51, 'mk': 0.549, 'pg': 0.523, 'sf': 0.271}

fns = ['dragon glass-Table 1.csv',
        'Joel\'s gravel-Table 1.csv',
        'marble chips-Table 1.csv',
        'Michael\'s gravel-Table 1.csv',
        'pea gravel-Table 1.csv',
        'shell fragments-Table 1.csv']

grain_name = {'dragon glass-Table 1.csv': 'dn',
        'Joel\'s gravel-Table 1.csv': 'jo',
        'marble chips-Table 1.csv': 'mb',
        'Michael\'s gravel-Table 1.csv': 'mk',
        'pea gravel-Table 1.csv': 'pg',
        'shell fragments-Table 1.csv': 'sf'}

df_shape_taylor = {}
for fn in fns:
    temp = pd.read_csv('2022_taylor_shape_data/'+fn)
    df_shape_taylor[grain_name[fn]] = {'A': temp['a (long)'].mean(), 'B': temp['b (mid)'].mean(), 'C': temp['c (short)'].mean(), 'CSF': temp['CSF'].mean()}

df_shape_taylor = pd.DataFrame(df_shape_taylor).T
A = {grain: df_shape_taylor.A[grain]/1000 for grain in grains}
# dA = {grain: df_shape_taylor.dA[grain] for grain in grains}
B = {grain: df_shape_taylor.B[grain]/1000 for grain in grains}
# dA = {grain: df_shape_taylor.dB[grain] for grain in grains}
C = {grain: df_shape_taylor.C[grain]/1000 for grain in grains}
# dA = {grain: df_shape_taylor.dC[grain] for grain in grains}
CSF = {grain: df_shape_taylor.CSF[grain] for grain in grains}
# dA = {grain: df_shape_taylor.dCSF[grain] for grain in grains}

# %%
# calculate cstar and mustar

bed_angle = 3.5
aor_gb = 24.736
cstar, mustar = {}, {}
for grain in grains:
    cstar[grain] = ((vel_VES[grain])/(vel[grain]))**2 * CSF[grain]
    mugb = (np.tan(np.radians(aor_gb)) - np.tan(np.radians(bed_angle)))
    mustar[grain] = (np.tan(np.radians(aor[grain])) - np.tan(np.radians(bed_angle)))/mugb

# %%
database = {'AoR': aor,
        'dAoR': daor,
        'rho': rho,
        # 'drho': drho,
        'mu*': mustar,
        # 'dmu*': dmustar,
        'C*': cstar,
        # 'dC*': dcstar,
        'C*/mu*': {grain: cstar[grain]/mustar[grain] for grain in grains},
        'set_vel_meas': vel,
        'dset_vel_meas': dvel,
        'set_vel_VES': vel_VES,
        'dset_vel_VES': dvel_VES,
        'CD/Co': {grain: (vel_VES[grain]/vel[grain])**2 for grain in grains},
        'CSF': CSF,
        # 'dCSF': dCSF,
        'A': A,
        # 'dA': dA,
        'B': B,
        # 'dB': dB,
        'C': C,
        # 'dC': dC,
        'Dn': d_VES,
        # 'dDn': dd_VES,
        }

df_all = pd.DataFrame(database)
df_all.to_csv('grain_database_2022.csv')

# %%
df_all
# %%
